# _*_ coding: utf-8 _*_
# 导入所需的库（仅保留评估必要的库）
import os
import random
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


# 设置随机种子以保证结果可重复性（保留，确保评估结果一致）
def setup_seed(seed):
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True


# 初始化随机种子
setup_seed(0)

# 设备配置（保留，自动选择GPU/CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 数据预处理（仅保留验证集的处理，删除训练集增强）
transform_val = transforms.Compose([
    transforms.Resize((224, 224)),  # 固定大小，无随机操作
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 与训练时保持一致
])

# 加载验证集数据（删除训练集加载代码）
val_dir = '../dataset/flower_data/val'
val_dataset = datasets.ImageFolder(
    val_dir,
    transform=transform_val
)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)  # 验证集不打乱


# 模型结构定义（必须保留，且与训练时完全一致）
class AlexNet(nn.Module):
    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 96, 11, 4, 2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(96, 256, 5, 1, 2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(256, 384, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


# 加载预训练模型（核心修改：移除训练代码，直接加载模型）
model_path = '../model/AlexNet_5.pth'  # 模型文件路径
model = AlexNet(num_classes=5).to(device)  # 初始化模型（类别数与训练时一致）
model.load_state_dict(torch.load(model_path, map_location=device))  # 加载权重
model.eval()  # 强制设置为评估模式（关闭dropout等训练特有的层）

# 模型评估（保留并简化评估逻辑）
correct = 0
total = 0
predicted_labels = []
true_labels = []

with torch.no_grad():  # 禁用梯度计算，节省内存并加速
    for images, labels in val_loader:
        images = images.to(device)
        labels = labels.to(device)

        # 仅执行前向传播（无反向传播）
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)

        # 统计结果
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        predicted_labels.extend(predicted.cpu().numpy())  # 转移到CPU并记录
        true_labels.extend(labels.cpu().numpy())

# 输出评估指标
print(f'模型在验证集上的准确率: {100 * correct / total:.4f} %')


